import torch
import torch.nn.functional as F
import numpy as np
from sklearn.metrics import roc_auc_score, f1_score

# 宏观F1分数
def eval_f1(y_true, y_pred):
    # 将 PyTorch 张量 y_true 从计算图中分离，移动到 CPU，并转换为 NumPy 数组。
    y_true = y_true.detach().cpu().numpy()
    # 对预测张量 y_pred 进行 argmax 操作，找到最大概率的类别索引。
    # 将结果从计算图中分离，移动到 CPU，并转换为 NumPy 数组。
    y_pred = y_pred.argmax(dim=-1, keepdim=True).detach().cpu().numpy()
    # 使用 f1_score 计算宏观平均 F1 分数。
    f1 = f1_score(y_true, y_pred, average='macro')
    return f1

# 准确率
def eval_acc(y_true, y_pred):
    # 初始化准确率列表，用于存储每个类别的准确率。
    acc_list = []
    # 将 PyTorch 张量 y_true 从计算图中分离，移动到 CPU，并转换为 NumPy 数组。
    y_true = y_true.detach().cpu().numpy()
    # 对预测张量 y_pred 进行 argmax 操作，找到最大概率的类别索引。将结果从计算图中分离，移动到 CPU，并转换为 NumPy 数组。
    y_pred = y_pred.argmax(dim=-1, keepdim=True).detach().cpu().numpy()

    # 遍历每个类别
    for i in range(y_true.shape[1]):
        # 这一行代码的目的是生成一个布尔掩码，表示哪些标签是有效的。通常在处理缺失标签或未标注数据时使用这种方法。
        # y_true[:, i] == y_true[:, i] 是逐元素比较操作，检查每个元素是否等于自身。对于大多数情况，这个比较总是返回 True，除非元素是 NaN（Not a Number）。因为 NaN 不等于自身。
        is_labeled = y_true[:, i] == y_true[:, i]
        # 这一行代码的目的是生成一个布尔数组，用于表示哪些预测值与真实标签相匹配（即预测正确）。
        # 筛选有效标签和预测：
        # y_true[is_labeled, i]：根据布尔掩码 is_labeled 筛选出有效的真实标签。
        # y_pred[is_labeled, i]：根据布尔掩码 is_labeled 筛选出有效的预测值。
        # 逐元素比较：
        # == 运算符用于逐元素比较 y_true 和 y_pred 中的对应元素，生成一个布尔数组 correct，表示哪些预测值与真实标签相匹配。
        correct = y_true[is_labeled, i] == y_pred[is_labeled, i]
        # 计算正确预测的比例，并添加到准确率列表中。
        acc_list.append(float(np.sum(correct))/len(correct))

    # 计算所有类别准确率的平均值，并返回结果。
    return sum(acc_list)/len(acc_list)


def eval_rocauc(y_true, y_pred):
    """ adapted from ogb
    https://github.com/snap-stanford/ogb/blob/master/ogb/nodeproppred/evaluate.py"""
    # 初始化 ROC-AUC 列表，用于存储每个类别的 ROC-AUC 分数。
    rocauc_list = []
    # 将 PyTorch 张量 y_true 从计算图中分离，移动到 CPU，并转换为 NumPy 数组。
    y_true = y_true.detach().cpu().numpy()
    if y_true.shape[1] == 1:
        # 对于单类分类：
        # 对预测值 y_pred 进行 softmax 操作，取第二个类别的概率。
        # 将结果移动到 CPU，并转换为 NumPy 数组。
        # use the predicted class for single-class classification
        # 这一行代码的目的是对预测值 y_pred 进行 softmax 操作，然后提取第二类的概率（假设这是二分类问题中的正类），再调整维度，并将结果转换为 NumPy 数组。
        # F.softmax(y_pred, dim=-1)：对 y_pred 的最后一个维度（类别维度）进行 softmax 操作。
        # [:, 1]：提取 softmax 结果中第二类的概率（假设这是二分类问题中的正类）。
        # unsqueeze(1)：在第 1 维度（列维度）插入一个新的维度，使得结果形状变为 [N, 1]。
        y_pred = F.softmax(y_pred, dim=-1)[:, 1].unsqueeze(1).cpu().numpy()
    else:
        # 对于多类分类：
        # 将预测张量 y_pred 从计算图中分离，移动到 CPU，并转换为 NumPy 数组。
        y_pred = y_pred.detach().cpu().numpy()

    # 遍历每个类别
    for i in range(y_true.shape[1]):
        # AUC is only defined when there is at least one positive data.
        # 这一段代码的目的是计算特定类别 i 的 ROC-AUC 分数。为了确保 ROC-AUC 分数的有效性，代码首先检查该类别是否有至少一个正样本和一个负样本，然后计算该类别的 ROC-AUC 分数。
        # np.sum(y_true[:, i] == 1) > 0：计算类别 i 的标签中值为 1 的数量，并检查是否大于 0。
        # np.sum(y_true[:, i] == 0) > 0：计算类别 i 的标签中值为 0 的数量，并检查是否大于 0。
        # 只有在类别 i 的标签中同时存在正样本（值为 1）和负样本（值为 0）时，才继续进行计算，否则跳过该类别的计算。
        if np.sum(y_true[:, i] == 1) > 0 and np.sum(y_true[:, i] == 0) > 0:
            # y_true[:, i] == y_true[:, i]：逐元素比较 y_true[:, i] 与自身，生成布尔掩码 is_labeled。
            # 对于大多数情况，这个比较总是返回 True，除非元素是 NaN（因为 NaN 不等于自身）。
            is_labeled = y_true[:, i] == y_true[:, i]
            # 计算 ROC-AUC 分数
            # y_true[is_labeled, i]：根据布尔掩码 is_labeled 筛选出有效的真实标签。
            # y_pred[is_labeled, i]：根据布尔掩码 is_labeled 筛选出有效的预测值。
            # roc_auc_score(y_true[is_labeled, i], y_pred[is_labeled, i])：计算筛选后的真实标签和预测值的 ROC-AUC 分数。
            score = roc_auc_score(y_true[is_labeled, i], y_pred[is_labeled, i])
            # 添加到列表中。
            rocauc_list.append(score)

    # 如果没有正样本，则引发错误。
    if len(rocauc_list) == 0:
        raise RuntimeError(
            'No positively labeled data available. Cannot compute ROC-AUC.')

    # 计算所有类别 ROC-AUC 分数的平均值，并返回结果。
    return sum(rocauc_list)/len(rocauc_list)

# @torch.no_grad()：禁用梯度计算，以提高评估过程的速度和节省内存。
@torch.no_grad()
def evaluate_full(model, dataset, eval_func):
    # 设置模型为评估模式，影响 Dropout 和 BatchNorm 等层的行为。
    model.eval()
    # 获取训练、验证、测试和 OOD 测试集的索引。
    train_idx, valid_idx, test_in_idx, test_ood_idx = dataset.train_idx, dataset.valid_idx, dataset.test_in_idx, dataset.test_ood_idx
    # 将标签张量移动到 CPU。
    y = dataset.y.cpu()
    # 使用模型对数据集进行前向传播，获取预测结果，并将结果移动到 CPU。
    out = model(dataset.x, dataset.edge_index).cpu()

    # 使用评估函数计算训练集/验证集/测试集的准确率。
    train_acc = eval_func(y[train_idx], out[train_idx])
    valid_acc = eval_func(y[valid_idx], out[valid_idx])
    test_in_acc = eval_func(y[test_in_idx], out[test_in_idx])
    # 初始化 OOD 测试集准确率列表。
    test_ood_accs = []
    # 遍历每个 OOD 测试集
    for t in test_ood_idx:
        # 使用评估函数计算每个 OOD 测试集的准确率，并添加到列表中。
        test_ood_accs.append(eval_func(y[t], out[t]))
    # 将训练、验证、测试和 OOD 测试集的准确率合并成结果列表。
    result = [train_acc, valid_acc, test_in_acc] + test_ood_accs
    # 返回结果列表
    return result
